import torch
import torch.nn as nn
class 丰度预测模型(nn.Module):
    # 注意：以下所有方法都要缩进，属于类的成员，，

    def __init__(self, input_dim=100, hidden_dim=256, output_dim=100):
        super(丰度预测模型, self).__init__()
        # 输入：s向量（100维），输出：p向量（100维）
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),  # 输入层到隐藏层
            nn.ReLU(),  # 激活函数
            nn.BatchNorm1d(hidden_dim),  # 批量归一化（加速训练）
            nn.Dropout(0.2),  # dropout（防止过拟合）
            nn.Linear(hidden_dim, hidden_dim // 2),  # 第二个隐藏层
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, output_dim),  # 输出层（100维丰度）
            nn.Softmax(dim=1)  # 输出转为概率分布（和为1）
        )

    # 关键修正：forward方法必须缩进，作为类的方法
    def forward(self, x):
        # x: 输入的s向量，形状：(batch_size, 100)
        return self.layers(x)